
from classfier.ModelUtil import ModelUtil
from utils.data_split import DataSplitUtil
import matplotlib.pyplot as plt
import pandas as pd
from scipy.stats import ttest_ind, mannwhitneyu
import numpy as np
from feature_extract.Lasso import LassoUtil
from feature_extract.forward_selection import ForwardSelectionUtil
from feature_extract.mRMR import MRMRUtil
from feature_extract.random_forest import RandomForestUtil
from feature_extract.rfe import RFEUtil
from utils.draw_pic import plot_mutil_model_roc
from utils.utils import *
import joblib
exclude_columns = ['Patient_ID', target_column]
clinic_select_features = [ 'Tumour Size(cm)', 'Age(years)', 'ER','PR','HER2']
# 用来暂存原始数据
data_dict = {}
# 用来暂存结果
result_dict = {}


def get_feature_extract_method(dp,srs,krs):
    method_class_map = {
        'lasso': LassoUtil,
        'rf': RandomForestUtil,
        'fs': ForwardSelectionUtil,
        'mrmr': MRMRUtil,
        'rfe': RFEUtil
    }
    temp_done_flag = opj(feature_select_path, 'done', str(srs),str(krs) + '_')
    for select_method in feature_select_method:
        feature_select_method_done_flag = temp_done_flag + select_method + '.txt'
        last_feature_select_method_done_flag = temp_done_flag + feature_select_method[-1] + '.txt'
        if ope(last_feature_select_method_done_flag):
            print(f"随机种子下{srs}，特征提取已完成，跳过。")
            break  # ✅ 已完成就跳过
        if ope(feature_select_method_done_flag):
            print(f"{select_method} 特征提取已完成")
            continue
        util_class = method_class_map.get(select_method)
        if util_class is None:
            print(f"未知方法：{select_method}，跳过")
            continue
        util_instance = util_class(dp, split_random_state=srs,pic_save_path=divide_result_path,selectK_random_state=krs)
        for data_type in data_type_list:
            df = pd.read_csv(opj(merge_data_path, data_type.replace('FeatureSummary', '') + '.csv'))
            flag = util_instance.extract_feature(df, data_type.replace('FeatureSummary', ''))

            if flag:
                print(f"跳过随机特征筛选种子 {krs} 的后续操作，因为方法{select_method}特征筛选{data_type.replace('FeatureSummary', '')}得到特征个数小于2")
                return False
        print(f"{select_method}特征提取已完成")
        util_instance.save_total_df()
        done(feature_select_method_done_flag)
    return True

def print_significance(result_df, alpha=0.001):
    # print(result_df)
    t_all_significant = (result_df['t_pvalue'] < alpha).all()
    u_all_significant = (result_df['u_pvalue'] < alpha).all()
    return ('T' if t_all_significant else 'F') + ('T' if u_all_significant else 'F')

def compare_features_by_label(X, y, feature_list):
    df = X.copy()
    df.columns = feature_list
    df[target_column] = y

    results = []

    for feature in feature_list:
        values_0 = df[df[target_column] == 0][feature]
        values_1 = df[df[target_column] == 1][feature]

        # 判断是否为正态分布可以用正态性检验，这里简化用 t 检验 + U检验都跑
        t_stat, t_p = ttest_ind(values_0, values_1, equal_var=False)
        u_stat, u_p = mannwhitneyu(values_0, values_1, alternative='two-sided')

        results.append({
            'feature': feature,
            'mean_y=0': np.mean(values_0),
            'mean_y=1': np.mean(values_1),
            't_pvalue': t_p,
            'u_pvalue': u_p
        })

    return pd.DataFrame(results)

def getData(selectK_random_state,split_random_state,feature_select_type):
    # 加载数据
    temp_dict = {}
    select_path = feature_path_map.get(feature_select_type)
    for t in data_to_solve_list[:2]:
        temp_dict[t] = pd.read_csv(opj(select_path, str(split_random_state),str(selectK_random_state), f"{t}.csv"))
    for t in data_to_solve_list[2:]:
        temp_dict[t] = pd.read_csv(opj(select_path, str(split_random_state),str(selectK_random_state), f"{t}.csv"))
    select_dict = {feature_select_type:temp_dict}
    data_dict[split_random_state] = select_dict

# 临床模型开发
def getClinicModel(split_random_state,clf,dp,model_random_state,krs):
    clinic_data = pd.read_csv(clinic_data_path)
    clinic_data = clinic_data[clinic_select_features + exclude_columns]
    sp = opj(divide_result_path, str(split_random_state),str(krs),str(model_random_state), 'metric', 'Clinic')
    clinic_model = ModelUtil(dp=dp, X=clinic_data, split_random_state=split_random_state, save_path=sp, data_type='Clinic', model_type=clf,model_random_state=model_random_state).get_model()
    path = opj(signature_model_weight_path,'Clinic')
    md(path)
    joblib.dump(clinic_model, opj(path,'model.pkl'))
    return clinic_model.getSignature()

# WSI使用三个独立模型的特征进一步训练模型
def getFinalModel_by_three_to_one(split_random_state,feature_select_type,clf,dp,selectK_random_state,model_random_state):
    model_dict = {}
    # 分别获取三个特征的独立模型
    for t in data_to_solve_list[:-1]:
        sp = opj(divide_result_path, str(split_random_state),str(selectK_random_state),str(model_random_state), feature_select_type + '+' + clf, t)
        model = ModelUtil(dp=dp, X=data_dict[split_random_state][feature_select_type][t], split_random_state=split_random_state,
                          save_path=sp, data_type=t, model_type=clf,model_random_state=model_random_state).get_model()
        model_dict[t] = model
        path = opj(signature_model_weight_path,t.replace('Roi','Nuclei'))
        md(path)
        joblib.dump(model,opj(path,'model.pkl'))
    # 三合一
    train_data = None
    test_data = None
    for t in data_to_solve_list[:3]:
        type_train_signature, type_test_signature = model_dict[t].getSignature()
        if train_data is None:
            train_data = type_train_signature
        else:
            train_data = train_data.merge(type_train_signature, left_index=True, right_index=True, how='inner')
        if test_data is None:
            test_data = type_test_signature
        else:
            test_data = test_data.merge(type_test_signature, left_index=True, right_index=True, how='inner')
    X = pd.concat([train_data, test_data], axis=0)
    X = X.reset_index()  # 这会把原索引变成一列并重置索引
    X.rename(columns={0: 'Patient_ID'}, inplace=True)  # 将索引列重命名为 Patient_ID
    temp_data = pd.read_csv(opj(merge_data_path, 'Roi.csv'))[['Patient_ID', target_column]]
    X = X.merge(temp_data, on='Patient_ID', how='inner').sort_values(by='Patient_ID', ascending=True).reset_index(
        drop=True)
    X_train, X_test, _, _, _, _ = dp.split_train_and_test(X,"")
    save_path =  opj(divide_result_path, str(split_random_state),str(selectK_random_state),str(model_random_state),feature_select_type+'+'+clf,'final')
    getFinalModel(X_train, X_test,split_random_state,save_path,clf,dp,model_random_state=model_random_state,selectK_random_state=selectK_random_state)

def getFinalModel(wsi_train_signature,wsi_test_signature,split_random_state,sp,clf,dp,model_random_state,selectK_random_state):
    clinic_train_signature, clinic_test_signature = getClinicModel(split_random_state,clf,dp,model_random_state=model_random_state,krs=selectK_random_state)
    train_data = wsi_train_signature.merge(clinic_train_signature, left_index=True, right_index=True, how='inner')
    test_data = wsi_test_signature.merge(clinic_test_signature, left_index=True, right_index=True, how='inner')
    X = pd.concat([train_data, test_data], axis=0)
    X = X.reset_index()  # 这会把原索引变成一列并重置索引
    X.rename(columns={'index': 'Patient_ID'}, inplace=True)  # 将索引列重命名为 Patient_ID
    temp_data = pd.read_csv(opj(merge_data_path, 'Roi.csv'))[['Patient_ID', target_column]]
    X = X.merge(temp_data, on='Patient_ID', how='inner').sort_values(by='Patient_ID', ascending=True).reset_index(
        drop=True)
    X = X.rename(columns={"Roi": "Nuclei"})
    X.to_csv(signature_score_csv_path,index=False)
    model = ModelUtil(dp=dp,X=X,split_random_state=split_random_state,save_path=sp,data_type='total',model_type=clf,model_random_state=model_random_state).get_model()


def main():
    sys.stdout = Logger(opj(base_path, 'logs', str(datetime.date.today()) + '_total.txt'))
    print('任务开始')
    for srs in split_random_state_list:
        print('-----------------------------------------------------------------------------')
        bp = opj(divide_result_path, str(srs))
        md(bp)
        print(f"🟩在分割随机种子{srs}下")
        # 划分数据
        data = pd.read_csv(opj(base_path, 'data','merge_data','BCNB','Collagen.csv'))
        data_split = DataSplitUtil(df=data, split_random_state=srs)
        data_split.test_split()
        print('数据集划分成功')
        for krs in selectK_random_state_list:
            # 特征提取
            print(f"🟦在特征选取随机种子 {krs} 下")
            print("特征提取流程开始")
            if not get_feature_extract_method(data_split,srs,krs):
                print("特征提取流程结束")
                continue  # 跳过该随机种子的后续模型构建和绘图
            print("特征提取流程结束")

            for mrs in model_random_state_list:
                print(f"🟥在模型随机种子 {mrs} 下")
                for fs_type, clf_type in combinations:
                    print(f'当前组合：特征选择={fs_type}, 分类器={clf_type}')
                    getData(krs, srs, fs_type)
                    getFinalModel_by_three_to_one(srs, fs_type, clf_type, data_split, krs, mrs)
                print(f"🟥模型随机种子 {mrs} 已完成")
            print(f"🟦特征选取随机种子 {krs} 已完成")

        print('-----------------------------------------------------------------------------')

if __name__ == '__main__':
   md(signature_model_weight_path)
   main()